import numpy as np
import matplotlib.pyplot as plt

# 定义材料和曲梁几何属性
E = 210e9  # 杨氏模量 (Pa)
Ixx = 8.3e-6  # 关于x轴的截面惯性矩 (m^4)
Iyy = 8.3e-6  # 关于y轴的截面惯性矩 (m^4)
J =16.6e-6  # 截面的极惯性矩 (m^4)
A = 0.01  # 截面面积 (m^2)
R = 1.0  # 曲梁的弯曲半径 (m)
theta = 2 * np.pi * 3 / 4 # 曲梁的角度 (两圈螺旋)
h = 0  # 螺旋高度 (m)

# 节点和单元定义
num_nodes = 20

# 节点数量
nodes_angle = np.linspace(0, theta, num_nodes)  # 在螺旋梁上均匀分布角度
nodes_height = np.linspace(0, h, num_nodes)  # 螺旋的高度变化
nodes = np.array([[R * np.cos(ang), R * np.sin(ang), z] for ang, z in zip(nodes_angle, nodes_height)])  # 曲梁节点坐标
num_elements = num_nodes - 1

# 自由度和单元定义
dof_per_node = 6  # 每个节点的自由度数量
elements = np.array([[i, i + 1] for i in range(num_elements)])

# 初始化全局刚度矩阵和力向量
K = np.zeros((num_nodes * dof_per_node, num_nodes * dof_per_node))  # 全局刚度矩阵
F = np.zeros(num_nodes * dof_per_node)  # 力向量


def beam_element_stiffness_curved_3d(E, Ixx, Iyy, J, A, L_elem):
    k_local = np.zeros((12, 12))

    EIxx = E * Ixx
    EIyy = E * Iyy
    GJ = E * J / (2 * (1 + 0.3))  # 假设泊松比为0.3
    EA = E * A

    k_local[0, 0] = k_local[6, 6] = EA / L_elem
    k_local[0, 6] = k_local[6, 0] = -EA / L_elem

    k_local[1, 1] = k_local[7, 7] = 12 * EIyy / L_elem ** 3
    k_local[1, 5] = k_local[5, 1] = 6 * EIyy / L_elem ** 2
    k_local[1, 7] = k_local[7, 1] = -12 * EIyy / L_elem ** 3
    k_local[1, 11] = k_local[11, 1] = 6 * EIyy / L_elem ** 2

    k_local[2, 2] = k_local[8, 8] = 12 * EIxx / L_elem ** 3
    k_local[2, 4] = k_local[4, 2] = -6 * EIxx / L_elem ** 2
    k_local[2, 8] = k_local[8, 2] = -12 * EIxx / L_elem ** 3
    k_local[2, 10] = k_local[10, 2] = -6 * EIxx / L_elem ** 2

    k_local[3, 3] = k_local[9, 9] = GJ / L_elem
    k_local[3, 9] = k_local[9, 3] = -GJ / L_elem

    k_local[4, 4] = k_local[10, 10] = 4 * EIyy / L_elem
    k_local[4, 8] = k_local[8, 4] = 6 * EIxx / L_elem ** 2
    k_local[4, 10] = k_local[10, 4] = 2 * EIyy / L_elem

    k_local[5, 5] = k_local[11, 11] = 4 * EIxx / L_elem
    k_local[5, 11] = k_local[11, 5] = 2 * EIxx / L_elem

    k_local[5, 7] = k_local[7, 5] = -6 * EIxx / L_elem ** 2
    k_local[7, 11] = k_local[11, 7] = -6 * EIxx / L_elem ** 2
    k_local[8, 10] = k_local[10, 8] = 6 * EIxx / L_elem ** 2

    print("单刚", k_local)
    return k_local


def assemble_global_stiffness(K, k_local, element):
    """
    汇总单元刚度矩阵到全局刚度矩阵
    """
    node1, node2 = element
    global_dof = np.array([node1 * dof_per_node + i for i in range(dof_per_node)] +
                          [node2 * dof_per_node + i for i in range(dof_per_node)])

    for i in range(12):
        for j in range(12):
            K[global_dof[i], global_dof[j]] += k_local[i, j]


# 汇总刚度矩阵
for element in elements:
    node1, node2 = element
    L_elem = R * (nodes_angle[node2] - nodes_angle[node1])  # 弧长 = 半径 * 角度差
    k_local = beam_element_stiffness_curved_3d(E, Ixx, Iyy, J, A, L_elem, R)
    assemble_global_stiffness(K, k_local, element)


# 应用边界条件和求解
def apply_boundary_conditions(K, F, boundary_conditions):
    """
    边界条件函数，消除约束自由度
    """
    for bc in boundary_conditions:
        dof = bc['dof']
        value = bc['value']
        K[dof, :] = 0
        K[:, dof] = 0
        K[dof, dof] = 1
        F[dof] = value


# 固定左端 (节点0的6个自由度)
boundary_conditions = [{'dof': i, 'value': 0.0} for i in range(dof_per_node)]
apply_boundary_conditions(K, F, boundary_conditions)

# 力向量（右端施加一个集中力，例如力作用在节点最后一个节点的y方向。倒数第五个自由度即最后一个节点的y方向）
F[-5] = 1000  # 节点50的Y方向力 (集中力 1000N)

# 求解位移
displacements = np.linalg.solve(K, F)

# 输出节点位移
print("节点位移：")
for i in range(num_nodes):
    disp = displacements[i * dof_per_node:(i + 1) * dof_per_node]
    print(f"节点 {i} 的位移: {disp[:3]}，旋转: {disp[3:]}")

# 可视化三维曲梁螺旋模型
def plot_helix(nodes):
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # 获取节点坐标
    x = nodes[:, 0]
    y = nodes[:, 1]
    z = nodes[:, 2]

    # 绘制螺旋曲梁
    ax.plot(x, y, z, '-o', label='Helix Beam', markersize=4)

    # 设置图形标签
    ax.set_xlabel('X 轴')
    ax.set_ylabel('Y 轴')
    ax.set_zlabel('Z 轴')
    ax.set_title('三维曲梁螺旋模型')
    ax.legend()

    plt.rcParams['font.sans-serif'] = ['SimHei']  # 黑体
    plt.rcParams['font.family'] = 'sans-serif'
    plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号

    plt.show()

# 绘制螺旋曲梁模型
plot_helix(nodes)